# Contact PM Aronow for questions (peter.aronow@yale.edu)
# Special thanks to Dean Eckles for extraordinarily helpful contributions

library(minqa)
library(nlme)

tsls.wfit <- function(X, Y, Z, weights, ...) {
  fs <- lm.wfit(Z, X, weights)
  XZ <- as.matrix(fs$fitted.values)
  colnames(XZ) <- colnames(X)
  ss <- lm.wfit(XZ, Y, weights)
  ss
}

complier.score <- function(D, Z, W, weights = NULL, link = pnorm, genoud = TRUE, num.iter = ifelse(genoud, 200, 1e4),onesided=FALSE) {
  if(!all.equal(sort(unique(c(D,0,1))), 0:1)) stop("Treatment must be binary.")
  if(!all.equal(sort(unique(c(Z,0,1))), 0:1)) stop("Assignment must be binary.")
  W <- as.matrix(W)
  if(length(D) != length(Z) || length(D) != nrow(W)) stop("Number of observations must be equal for all variables.")
  n.obs <- length(D)
  if (is.null(weights))
    weights <- rep(1, n.obs)
  
  variances <- apply(W,2,var)
  if (min(variances) == 0) stop("One or more predictors are constants.")

  complier.neg.logLik.fnc <- function (par) {
  	
    theta <- par[1:(length(par)/2)]
    theta.A <- par[(length(par)/2+1):(length(par))]
    
    p.AC <- link(rowSums(t(t(W1) * theta)))
    p.A <- link(rowSums(t(t(W1) * theta.A)))
    p.C <- p.AC - p.A*p.AC
    
    p.D <- p.AC * (1 - p.A) * Z + p.AC * p.A
    p.D <- pmax(pmin(p.D, 1 - 1e-5), 1e-5)
    
    return(
	-sum(weights*log((p.D)^D * (1 - p.D)^(1 - D))))
	  }
  
# one sided

	if (onesided == TRUE)   complier.neg.logLik.fnc <- function (par) {
    theta <- par[1:(length(par))]
    p.C <- link(rowSums(t(t(W1) * theta)))
    p.D <- p.C * Z
    return(
	-sum(log((p.D)^D * (1 - p.D)^(1 - D) * weights)))
	
	}

  # cheap starting values
  df <- data.frame(W = W)
 
  # normalize all inputs
  
  W <- apply(cbind(W),2,function(x) (x - mean(x))/sd(x))
  
if (onesided == FALSE) {  
  suppressWarnings(glm.CA <- glm(D ~ W, df, subset = (Z == 1), family = binomial(link="probit")))
  suppressWarnings(glm.A <- glm(D ~ W, df, subset = (Z == 0), family = binomial(link="probit")))
  CA.score.cheap <- predict(glm.CA, df, weights = weights, type = "response")
  A.score.cheap <- predict(glm.A, df, weights = weights, type = "response")
  p.A.given.AC <- pmin(pmax(A.score.cheap / CA.score.cheap, 0), 1)
  suppressWarnings(glm.A.given.AC <- lm(p.A.given.AC ~ W, weights = weights))
  starting <- c(glm.CA$coefficients, glm.A.given.AC$coefficients)
} else starting <- glm(D ~ W, df, subset = (Z == 1), family = binomial(link="probit"))$coefficients


  W1 <- cbind(1,W)
  
if (onesided == FALSE) {
  # optimize
    starting[starting < -4.5] <- -4.5
  	starting[starting > 4.5] <- 4.5
  	starting[!is.finite(starting)] <- 0

  if (genoud == FALSE) {
    opt.log <- bobyqa(starting, complier.neg.logLik.fnc,control=list(iprint=0,maxfun=num.iter),lower=-5,upper=5)
  } else {
    if(!isTRUE(require(rgenoud))) stop("rgenoud package required to use genoud = TRUE option.")
    opt.log <- genoud(complier.neg.logLik.fnc, (ncol(W1) * 2), starting.values = starting, pop.size = 1000, Domains = 5, max.generations = num.iter, hessian = FALSE, print.level = 2, hard.generation.limit = TRUE, gradient.check = FALSE ,optim.method = "Nelder-Mead")
  }
  CA.score <- (pnorm(rowSums(t(t(W1) * opt.log$par[1:(length(opt.log$par)/2)]))))
  A.score <- (pnorm(rowSums(t(t(W1) * opt.log$par[(length(opt.log$par)/2+1):(length(opt.log$par))])))) * CA.score
  C.score <- CA.score - A.score
  list(C.score = C.score, A.score = A.score, opt.log = opt.log)
} else {
	
	if (genoud == FALSE) {
    opt.log <- bobyqa(starting, complier.neg.logLik.fnc,control=list(iprint=2,maxfun=num.iter))
  } else {
    if(!isTRUE(require(rgenoud))) stop("rgenoud package required to use genoud = TRUE option.")
    opt.log <- genoud(complier.neg.logLik.fnc, (ncol(W1)), starting.values = starting, pop.size = 1000, Domains = 5, max.generations = num.iter, hessian = FALSE, print.level = 2, hard.generation.limit = TRUE, gradient.check = FALSE ,optim.method = "Nelder-Mead")
  }
  C.score <- (pnorm(rowSums(t(t(W1) * opt.log$par[1:(length(opt.log$par))]))))
  list(C.score = C.score, A.score = rep(0,n.obs), opt.log = opt.log)
}
	}


windsorize.probs <- function(prob, min.prob) {
if (!is.null(min.prob)) {
    if(min(prob) < min.prob) {
      warning("Small probabilities encountered. Replacing with min.prob.")
      prob[prob < min.prob] <- min.prob
    }
  }
  if (min(prob) <= 0 & is.null(min.prob)) {
    warning("Probabilities scores of zero encountered and min.prob not set. Replacing with smallest value > 0.")
    prob[prob == 0] <- min(prob[prob > 0])
  } 
  return(prob)
}
